from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
import numpy as np
from typing import Dict, Any
import os
import time
from datetime import datetime

class MLTrainer:
    def __init__(self):
        self.model = None
        self.training_status = "idle"
        self.history = {
            "loss": [],
            "accuracy": [],
            "timestamps": []
        }
        self.log_file = None
        
    def _create_log_file(self):
        """创建日志文件"""
        os.makedirs("app/logs", exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.log_file = f"app/logs/training_{timestamp}.log"
        return self.log_file
    
    def _log_metrics(self, loss: float, accuracy: float):
        """记录训练指标"""
        timestamp = time.time()
        log_entry = f"Validation Loss: {loss:.4f}, Accuracy: {accuracy:.2f}%"
        
        # 保存到历史记录
        self.history["loss"].append(loss)
        self.history["accuracy"].append(accuracy)
        self.history["timestamps"].append(timestamp)
        
        # 写入日志文件
        with open(self.log_file, "a") as f:
            f.write(f"{log_entry}\n")
            
        return log_entry
        
    async def start_training(self, params: Dict[str, Any]):
        """启动模型训练"""
        try:
            self.training_status = "training"
            self._create_log_file()
            
            # 生成示例数据
            X = np.random.rand(1000, 4)
            y = np.random.randint(0, 2, 1000)
            
            # 分割训练集和验证集
            X_train, X_val, y_train, y_val = train_test_split(
                X, y, test_size=0.2, random_state=42
            )
            
            self.model = LogisticRegression(**params)
            
            # 训练模型并记录指标
            self.model.fit(X_train, y_train)
            
            # 计算验证集上的指标
            y_pred = self.model.predict(X_val)
            y_pred_proba = self.model.predict_proba(X_val)
            
            # 计算交叉熵损失
            loss = -np.mean(
                [np.log(proba[true]) for proba, true in zip(y_pred_proba, y_val)]
            )
            
            # 计算准确率
            accuracy = (y_pred == y_val).mean() * 100
            
            # 记录指标
            self._log_metrics(loss, accuracy)
            
            self.training_status = "completed"
            return {
                "status": "success",
                "message": "Training completed successfully",
                "final_metrics": {
                    "loss": loss,
                    "accuracy": accuracy
                }
            }
            
        except Exception as e:
            self.training_status = "failed"
            return {"status": "error", "message": str(e)}
    
    def get_status(self):
        """获取训练状态"""
        return {
            "status": self.training_status,
            "model_info": str(self.model) if self.model else None,
            "history": self.history if self.history["loss"] else None,
            "log_file": self.log_file
        } 